#include "GDICapture.h"
#include "Utils.h"
#include <iostream>
#include <string>

GDICapture::GDICapture() :
    m_initialized(false),
    m_screenDC(NULL),
    m_memoryDC(NULL),
    m_currentDesktop(NULL),
    m_screenWidth(0),
    m_screenHeight(0) {
}

GDICapture::~GDICapture() {
    Cleanup();
}

bool GDICapture::Initialize() {
    Utils::DebugLog("GDICapture: Initializing...");

    m_screenWidth = GetSystemMetrics(SM_CXSCREEN);
    m_screenHeight = GetSystemMetrics(SM_CYSCREEN);

    if (m_screenWidth <= 0 || m_screenHeight <= 0) {
        m_lastError = "Invalid screen dimensions";
        return false;
    }

    HDESK hDesktop = GetThreadDesktop(GetCurrentThreadId());
    if (!hDesktop) {
        m_lastError = "Failed to get current desktop: " + Utils::GetLastWin32Error();
        return false;
    }

    if (!SetupDCsForDesktop(hDesktop)) {
        return false;
    }

    m_currentDesktop = hDesktop;
    m_initialized = true;

    Utils::DebugLog("GDICapture: Initialized successfully");
    return true;
}

void GDICapture::Cleanup() {
    Utils::DebugLog("GDICapture: Cleaning up...");

    CleanupDCs();

    if (m_currentDesktop) {
        m_currentDesktop = NULL;
    }

    m_initialized = false;
}

bool GDICapture::SetupDCsForDesktop(HDESK hDesktop) {
    CleanupDCs();

    HDESK hOldDesktop = GetThreadDesktop(GetCurrentThreadId());

    if (!SetThreadDesktop(hDesktop)) {
        Utils::DebugLog("GDICapture: Failed to set thread desktop, trying direct access");

        m_screenDC = GetDC(NULL);
        if (m_screenDC) {
            m_memoryDC = CreateCompatibleDC(m_screenDC);
            if (m_memoryDC) {
                return true;
            }
            ReleaseDC(NULL, m_screenDC);
            m_screenDC = NULL;
        }

        m_lastError = "Failed to set thread desktop: " + Utils::GetLastWin32Error();
        return false;
    }

    m_screenDC = GetDC(NULL);
    if (!m_screenDC) {
        m_lastError = "Failed to get screen DC: " + Utils::GetLastWin32Error();
        SetThreadDesktop(hOldDesktop);
        return false;
    }

    m_memoryDC = CreateCompatibleDC(m_screenDC);
    if (!m_memoryDC) {
        m_lastError = "Failed to create memory DC: " + Utils::GetLastWin32Error();
        ReleaseDC(NULL, m_screenDC);
        m_screenDC = NULL;
        SetThreadDesktop(hOldDesktop);
        return false;
    }

    return true;
}

void GDICapture::CleanupDCs() {
    if (m_memoryDC) {
        DeleteDC(m_memoryDC);
        m_memoryDC = NULL;
    }

    if (m_screenDC) {
        ReleaseDC(NULL, m_screenDC);
        m_screenDC = NULL;
    }
}

bool GDICapture::CaptureScreen(HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "GDI capture not initialized";
        return false;
    }

    // Always refresh DCs for secure desktop compatibility
    CleanupDCs();

    // Try multiple methods to get valid DCs
    HDESK hCurrentDesktop = GetThreadDesktop(GetCurrentThreadId());

    // Method 1: Try to switch to current desktop first
    if (!SetThreadDesktop(hCurrentDesktop)) {
        Utils::DebugLog("GDICapture: Cannot switch to current desktop, trying alternative method");
    }

    // Method 2: Try getting DC from desktop window
    HWND hDesktopWnd = GetDesktopWindow();
    m_screenDC = GetDC(hDesktopWnd);
    if (!m_screenDC) {
        // Method 3: Try getting DC from NULL (entire screen)
        m_screenDC = GetDC(NULL);
    }

    if (!m_screenDC) {
        m_lastError = "Failed to get any screen DC: " + Utils::GetLastWin32Error();
        return false;
    }

    m_memoryDC = CreateCompatibleDC(m_screenDC);
    if (!m_memoryDC) {
        m_lastError = "Failed to create memory DC: " + Utils::GetLastWin32Error();
        ReleaseDC(hDesktopWnd ? hDesktopWnd : NULL, m_screenDC);
        m_screenDC = NULL;
        return false;
    }

    m_currentDesktop = hCurrentDesktop;

    int currentWidth = GetSystemMetrics(SM_CXSCREEN);
    int currentHeight = GetSystemMetrics(SM_CYSCREEN);

    if (currentWidth != m_screenWidth || currentHeight != m_screenHeight) {
        m_screenWidth = currentWidth;
        m_screenHeight = currentHeight;
        Utils::DebugLog("GDICapture: Screen resolution changed");
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(m_screenDC, m_screenWidth, m_screenHeight);
    if (!hScreenBitmap) {
        m_lastError = "Failed to create compatible bitmap: " + Utils::GetLastWin32Error();
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(m_memoryDC, hScreenBitmap);
    if (!hOldBitmap || hOldBitmap == HGDI_ERROR) {
        m_lastError = "Failed to select bitmap into memory DC: " + Utils::GetLastWin32Error();
        DeleteObject(hScreenBitmap);
        return false;
    }

    BOOL result = FALSE;

    // Try different BitBlt methods for secure desktop compatibility

    // Method 1: Standard BitBlt
    result = BitBlt(m_memoryDC, 0, 0, m_screenWidth, m_screenHeight,
                   m_screenDC, 0, 0, SRCCOPY);

    if (!result) {
        Utils::DebugLog("GDICapture: Standard BitBlt failed, trying PrintWindow method");

        // Method 2: Try PrintWindow on desktop window
        HWND hDesktopWnd = GetDesktopWindow();
        if (hDesktopWnd) {
            result = PrintWindow(hDesktopWnd, m_memoryDC, PW_CLIENTONLY);
            if (!result) {
                result = PrintWindow(hDesktopWnd, m_memoryDC, 0);
            }
        }
    }

    if (!result) {
        Utils::DebugLog("GDICapture: PrintWindow failed, trying alternative BitBlt");

        // Method 3: Try with CAPTUREBLT flag for layered windows
        result = BitBlt(m_memoryDC, 0, 0, m_screenWidth, m_screenHeight,
                       m_screenDC, 0, 0, SRCCOPY | CAPTUREBLT);
    }

    SelectObject(m_memoryDC, hOldBitmap);

    if (!result) {
        Utils::DebugLog("GDICapture: All standard methods failed, trying alternative method");
        DeleteObject(hScreenBitmap);

        // Try alternative capture method for secure desktops
        return CaptureWithAlternativeMethod(hBitmap, width, height);
    }

    hBitmap = hScreenBitmap;
    width = m_screenWidth;
    height = m_screenHeight;

    return true;
}

bool GDICapture::CaptureWithAlternativeMethod(HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("GDICapture: Attempting alternative capture method for secure desktop");

    // Method: Try to enumerate all visible windows and composite them
    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);

    if (screenWidth <= 0 || screenHeight <= 0) {
        m_lastError = "Invalid screen dimensions for alternative method";
        return false;
    }

    // Create a bitmap for the screen
    HDC hScreenDC = GetDC(NULL);
    if (!hScreenDC) {
        m_lastError = "Failed to get screen DC for alternative method";
        return false;
    }

    HDC hMemDC = CreateCompatibleDC(hScreenDC);
    if (!hMemDC) {
        m_lastError = "Failed to create memory DC for alternative method";
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
    if (!hScreenBitmap) {
        m_lastError = "Failed to create bitmap for alternative method";
        DeleteDC(hMemDC);
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

    // Fill with black first
    RECT rect = { 0, 0, screenWidth, screenHeight };
    FillRect(hMemDC, &rect, (HBRUSH)GetStockObject(BLACK_BRUSH));

    // Try to get system color information
    COLORREF sysColor = GetSysColor(COLOR_DESKTOP);
    HBRUSH hBrush = CreateSolidBrush(sysColor);
    if (hBrush) {
        FillRect(hMemDC, &rect, hBrush);
        DeleteObject(hBrush);
    }

    // Try to capture visible windows
    HWND hWnd = GetTopWindow(GetDesktopWindow());
    while (hWnd) {
        if (IsWindowVisible(hWnd)) {
            RECT windowRect;
            if (GetWindowRect(hWnd, &windowRect)) {
                HDC hWndDC = GetWindowDC(hWnd);
                if (hWndDC) {
                    int wndWidth = windowRect.right - windowRect.left;
                    int wndHeight = windowRect.bottom - windowRect.top;

                    // Try PrintWindow first
                    if (!PrintWindow(hWnd, hMemDC, 0)) {
                        // Fallback to BitBlt from window DC
                        BitBlt(hMemDC, windowRect.left, windowRect.top, wndWidth, wndHeight,
                               hWndDC, 0, 0, SRCCOPY);
                    }

                    ReleaseDC(hWnd, hWndDC);
                }
            }
        }
        hWnd = GetNextWindow(hWnd, GW_HWNDNEXT);
    }

    SelectObject(hMemDC, hOldBitmap);
    DeleteDC(hMemDC);
    ReleaseDC(NULL, hScreenDC);

    hBitmap = hScreenBitmap;
    width = screenWidth;
    height = screenHeight;

    Utils::DebugLog("GDICapture: Alternative method completed");
    return true;
}